//===--- CSStep.h - Constraint Solver Steps -------------------------------===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2018 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//
//
// This file implements the \c SolverStep class and its related types,
// which is used by constraint solver to do iterative solving.
//
//===----------------------------------------------------------------------===//

#ifndef POLARPHP_SEMA_INTERNAL_CSSTEP_H
#define POLARPHP_SEMA_INTERNAL_CSSTEP_H

#include "polarphp/sema/internal/Constraint.h"
#include "polarphp/sema/internal/ConstraintGraph.h"
#include "polarphp/sema/internal/ConstraintSystem.h"
#include "polarphp/ast/Types.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/Optional.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/raw_ostream.h"
#include <memory>

using namespace llvm;

namespace polar {

namespace constraints {

class SolverStep;
class ComponentStep;

/// Represents available states which every
/// given step could be in during it's lifetime.
enum class StepState { Setup, Ready, Running, Suspended, Done };

/// Represents result of the step execution,
/// and can only be constructed by `SolverStep`.
struct StepResult {
   using Kind = ConstraintSystem::SolutionKind;

   friend class SolverStep;

private:
   Kind ResultKind;
   SmallVector<std::unique_ptr<SolverStep>, 4> NextSteps;

   StepResult(Kind kind) : ResultKind(kind) {}

   StepResult(Kind kind, std::unique_ptr<SolverStep> step) : ResultKind(kind) {
      NextSteps.push_back(std::move(step));
   }

   StepResult(Kind kind, SmallVectorImpl<std::unique_ptr<SolverStep>> &followup)
      : ResultKind(kind), NextSteps(std::move(followup)) {}

public:
   StepResult() = delete;

   Kind getKind() const { return ResultKind; }

   void transfer(SmallVectorImpl<std::unique_ptr<SolverStep>> &workList) {
      workList.reserve(NextSteps.size());
      for (auto &step : NextSteps)
         workList.push_back(std::move(step));
   }

private:
   static StepResult success() { return StepResult(Kind::Solved); }
   static StepResult failure() { return StepResult(Kind::Error); }

   static StepResult unsolved(std::unique_ptr<SolverStep> singleStep) {
      return StepResult(Kind::Unsolved, std::move(singleStep));
   }

   static StepResult
   unsolved(SmallVectorImpl<std::unique_ptr<SolverStep>> &followup) {
      return StepResult(Kind::Unsolved, followup);
   }
};

/// Represents a single independently solvable part of
/// the constraint system. And is a base class for all
/// different types of steps there are.
class SolverStep {
   friend class ConstraintSystem;

protected:
   ConstraintSystem &CS;

   StepState State = StepState::Setup;

   /// Once step is complete this is a container to hold finalized solutions.
   SmallVectorImpl<Solution> &Solutions;

public:
   explicit SolverStep(ConstraintSystem &cs,
                       SmallVectorImpl<Solution> &solutions)
      : CS(cs), Solutions(solutions) {}

   virtual ~SolverStep() {}

   /// \returns The current state of this step.
   StepState getState() const { return State; }

   /// Run preliminary setup (if needed) right
   /// before taking this step for the first time.
   virtual void setup() {}

   /// Try to move solver forward by simplifying constraints if possible.
   /// Such simplication might lead to either producing a solution, or
   /// creating a set of "follow-up" more granular steps to execute.
   ///
   /// \param prevFailed Indicate whether previous step
   ///        has failed (returned StepResult::Kind = Error),
   ///        this is useful to propagate failures when
   ///        unsolved steps are re-taken.
   ///
   /// \returns status and any follow-up steps to take before considering
   ///          this step solved or failed.
   virtual StepResult take(bool prevFailed) = 0;

   /// Try to resume previously suspended step.
   ///
   /// This happens after "follow-up" steps are done
   /// and all of the required information should be
   /// available to re-take this step.
   ///
   /// \param prevFailed Indicate whether previous step
   ///        has failed (returned StepResult::Kind = Error),
   ///        this is useful to propagate failures when
   ///        unsolved steps are re-taken.
   ///
   /// \returns status and any follow-up steps to take before considering
   ///          this step solved or failed.
   virtual StepResult resume(bool prevFailed) = 0;

   virtual void print(llvm::raw_ostream &Out) = 0;

protected:
   /// Transition this step into one of the available states.
   ///
   /// This is primarily driven by execution of the step itself and
   /// the solver, while it executes the work list.
   ///
   /// \param newState The new state this step should be in.
   void transitionTo(StepState newState) {
#ifndef NDEBUG
      // Make sure that ordering of the state transitions is correct,
      // because `setup -> ready -> running [-> suspended]* -> done`
      // is the only reasonable state transition path.
      switch (State) {
         case StepState::Setup:
            assert(newState == StepState::Ready);
            break;

         case StepState::Ready:
            assert(newState == StepState::Running);
            break;

         case StepState::Running:
            assert(newState == StepState::Suspended || newState == StepState::Done);
            break;

         case StepState::Suspended:
            assert(newState == StepState::Running);
            break;

         case StepState::Done:
            llvm_unreachable("step is already done.");
      }
#endif

      State = newState;
   }

   StepResult done(bool isSuccess) {
      transitionTo(StepState::Done);
      return isSuccess ? StepResult::success() : StepResult::failure();
   }

   StepResult replaceWith(std::unique_ptr<SolverStep> replacement) {
      transitionTo(StepState::Done);
      return StepResult(StepResult::Kind::Solved, std::move(replacement));
   }

   StepResult suspend(std::unique_ptr<SolverStep> followup) {
      transitionTo(StepState::Suspended);
      return StepResult::unsolved(std::move(followup));
   }

   StepResult suspend(SmallVectorImpl<std::unique_ptr<SolverStep>> &followup) {
      transitionTo(StepState::Suspended);
      return StepResult::unsolved(followup);
   }

   /// Erase constraint from the constraint system (include constraint graph)
   /// and return the constraint which follows it.
   ConstraintList::iterator erase(Constraint *constraint) {
      CS.CG.removeConstraint(constraint);
      return CS.InactiveConstraints.erase(constraint);
   }

   void restore(ConstraintList::iterator &iterator, Constraint *constraint) {
      CS.InactiveConstraints.insert(iterator, constraint);
      CS.CG.addConstraint(constraint);
   }

   ResolvedOverloadSetListItem *getResolvedOverloads() const {
      return CS.resolvedOverloadSets;
   }

   void recordDisjunctionChoice(ConstraintLocator *disjunctionLocator,
                                unsigned index) const {
      CS.recordDisjunctionChoice(disjunctionLocator, index);
   }

   Score getCurrentScore() const { return CS.CurrentScore; }

   Optional<Score> getBestScore() const { return CS.solverState->BestScore; }

   void filterSolutions(SmallVectorImpl<Solution> &solutions, bool minimize) {
      if (!CS.retainAllSolutions())
         CS.filterSolutions(solutions, minimize);
   }

   /// Check whether constraint solver is running in "debug" mode,
   /// which should output diagnostic information.
   bool isDebugMode() const {
      return CS.getAstContext().TypeCheckerOpts.DebugConstraintSolver;
   }

   llvm::raw_ostream &getDebugLogger(bool indent = true) const {
      auto &log = CS.getAstContext().TypeCheckerDebug->getStream();
      return indent ? log.indent(CS.solverState->depth * 2) : log;
   }
};

/// `SplitterStep` is responsible for running connected components
/// algorithm to determine how many independent sub-systems there are.
/// Once that's done it would create one `ComponentStep` per such
/// sub-system, and move to try to solve each and then merge partial
/// solutions produced by components into complete solution(s).
class SplitterStep final : public SolverStep {
   // Set of constraints associated with each component, after
   // component steps are complete, all of the constraints are
   // returned back to the work-list in their original order.
   SmallVector<ConstraintList, 4> Components;
   // Partial solutions associated with given step, each element
   // of the array presents a disjoint component (or follow-up step)
   // that current step has been split into.
   std::unique_ptr<SmallVector<Solution, 4>[]> PartialSolutions = nullptr;

   SmallVector<Constraint *, 4> OrphanedConstraints;

   /// Whether to include the partial results of this component in the final
   /// merged results.
   SmallVector<bool, 4> IncludeInMergedResults;

public:
   SplitterStep(ConstraintSystem &cs, SmallVectorImpl<Solution> &solutions)
      : SolverStep(cs, solutions) {}

   StepResult take(bool prevFailed) override;
   StepResult resume(bool prevFailed) override;

   void print(llvm::raw_ostream &Out) override {
      Out << "SplitterStep with #" << Components.size() << " components\n";
   }

private:
   /// If current step needs follow-up steps to get completely solved,
   /// let's compute them using connected components algorithm.
   void computeFollowupSteps(
      SmallVectorImpl<std::unique_ptr<SolverStep>> &steps);

   /// Once all of the follow-up steps are complete, let's try
   /// to merge resulting solutions together, to form final solution(s)
   /// for this step.
   ///
   /// \returns true if there are any solutions, false otherwise.
   bool mergePartialSolutions() const;
};

/// `DependentComponentSplitterStep` is responsible for composing the partial
/// solutions from other components (on which this component depends) into
/// the inputs based on which we can solve a particular component.
class DependentComponentSplitterStep final : public SolverStep {
   /// Constraints "in scope" of this step.
   ConstraintList *Constraints;

   /// Index into the parent splitter step.
   unsigned Index;

   /// The component that has dependencies.
   ConstraintGraph::Component Component;

   /// Array containing all of the partial solutions for the parent split.
   MutableArrayRef<SmallVector<Solution, 4>> AllPartialSolutions;

   /// Take all of the constraints in this component and put them into
   /// \c Constraints.
   void injectConstraints() {
      for (auto constraint : Component.getConstraints()) {
         Constraints->erase(constraint);
         Constraints->push_back(constraint);
      }
   }

public:
   DependentComponentSplitterStep(
      ConstraintSystem &cs,
      ConstraintList *constraints,
      unsigned index,
      ConstraintGraph::Component &&component,
      MutableArrayRef<SmallVector<Solution, 4>> allPartialSolutions)
      : SolverStep(cs, allPartialSolutions[index]), Constraints(constraints),
        Index(index), Component(std::move(component)),
        AllPartialSolutions(allPartialSolutions) {
      assert(!Component.dependsOn.empty() && "Should use ComponentStep");
      injectConstraints();
   }

   StepResult take(bool prevFailed) override;
   StepResult resume(bool prevFailed) override;

   void print(llvm::raw_ostream &Out) override;
};


/// `ComponentStep` represents a set of type variables and related
/// constraints which could be solved independently. It's further
/// simplified into "binding" steps which attempt type variable and
/// disjunction choices.
class ComponentStep final : public SolverStep {
   class Scope {
      ConstraintSystem &CS;
      ConstraintSystem::SolverScope *SolverScope;

      SetVector<TypeVariableType *> TypeVars;
      ConstraintSystem::SolverScope *PrevPartialScope = nullptr;

      // The component this scope is associated with.
      ComponentStep &Component;

   public:
      Scope(ComponentStep &component);

      ~Scope() {
         delete SolverScope; // rewind back all of the changes.
         CS.solverState->PartialSolutionScope = PrevPartialScope;

         // return all of the saved type variables back to the system.
         CS.TypeVariables = std::move(TypeVars);
         // return all of the saved constraints back to the component.
         auto &constraints = *Component.Constraints;
         constraints.splice(constraints.end(), CS.InactiveConstraints);
      }
   };

   /// The position of the component in the set of
   /// components produced by "split" step.
   unsigned Index;

   /// Indicates whether this is only component produced
   /// by "split" step. This information opens optimization
   /// opportunity, because if there are no other components,
   /// constraint system doesn't have to pruned from
   /// unrelated type variables and their constraints.
   bool IsSingle;

   /// The score associated with constraint system before
   /// the component step is taken.
   Score OriginalScore;

   /// The original best score computed before any of the
   /// component steps belonging to the same "split" are taken.
   Optional<Score> OriginalBestScore;

   /// If this step depends on other smaller steps to be solved first
   /// we need to keep active scope until all of the work is done.
   std::unique_ptr<Scope> ComponentScope = nullptr;

   /// Type variables and constraints "in scope" of this step.
   TinyPtrVector<TypeVariableType *> TypeVars;
   /// Constraints "in scope" of this step.
   ConstraintList *Constraints;

   /// The set of partial solutions that should be composed before evaluating
   /// this component.
   SmallVector<const Solution *, 2> DependsOnPartialSolutions;

   /// Constraint which doesn't have any free type variables associated
   /// with it, which makes it disconnected in the graph.
   Constraint *OrphanedConstraint = nullptr;

public:
   /// Create a single component step.
   ComponentStep(ConstraintSystem &cs, unsigned index,
                 ConstraintList *constraints,
                 SmallVectorImpl<Solution> &solutions)
      : SolverStep(cs, solutions), Index(index), IsSingle(true),
        OriginalScore(getCurrentScore()), OriginalBestScore(getBestScore()),
        Constraints(constraints) {}

   /// Create a component step from a constraint graph component.
   ComponentStep(ConstraintSystem &cs, unsigned index,
                 ConstraintList *constraints,
                 ConstraintGraph::Component &&component,
                 SmallVectorImpl<Solution> &solutions)
      : SolverStep(cs, solutions), Index(index), IsSingle(false),
        OriginalScore(getCurrentScore()), OriginalBestScore(getBestScore()),
        Constraints(constraints) {
      if (component.isOrphaned()) {
         assert(component.getConstraints().size() == 1);
         OrphanedConstraint = component.getConstraints().front();
      } else {
         assert(component.typeVars.size() > 0);
      }

      TypeVars = std::move(component.typeVars);

      for (auto constraint : component.getConstraints()) {
         constraints->erase(constraint);
         Constraints->push_back(constraint);
      }

      assert(component.dependsOn.empty());
   }

   /// Create a component step that composes existing partial solutions before
   /// solving constraints.
   ComponentStep(
      ConstraintSystem &cs, unsigned index,
      ConstraintList *constraints,
      const ConstraintGraph::Component &component,
      llvm::SmallVectorImpl<const Solution *> &&dependsOnPartialSolutions,
      SmallVectorImpl<Solution> &solutions)
      : SolverStep(cs, solutions), Index(index), IsSingle(false),
        OriginalScore(getCurrentScore()), OriginalBestScore(getBestScore()),
        Constraints(constraints),
        DependsOnPartialSolutions(std::move(dependsOnPartialSolutions)) {
      TypeVars = component.typeVars;
      assert(DependsOnPartialSolutions.size() == component.dependsOn.size());

      for (auto constraint : component.getConstraints()) {
         constraints->erase(constraint);
         Constraints->push_back(constraint);
      }
   }

   StepResult take(bool prevFailed) override;

   StepResult resume(bool prevFailed) override { return finalize(!prevFailed); }

   void print(llvm::raw_ostream &Out) override {
      Out << "ComponentStep with at #" << Index << '\n';
   }

private:
   void setupScope() {
      // If this is a single component, there is no need
      // to preliminary modify constraint system or log anything.
      if (IsSingle)
         return;

      if (isDebugMode())
         getDebugLogger() << "(solving component #" << Index << '\n';

      ComponentScope = std::make_unique<Scope>(*this);

      // If this component has orphaned constraint attached,
      // let's return it to the graph.
      CS.CG.setOrphanedConstraint(OrphanedConstraint);
   }

   /// Finalize current component by either cleanup if sub-tasks
   /// have failed, or solution generation and minimization.
   StepResult finalize(bool isSuccess);
};

template <typename P> class BindingStep : public SolverStep {
   using Scope = ConstraintSystem::SolverScope;

   P Producer;

protected:
   /// Indicates whether any of the attempted bindings
   /// produced a solution.
   bool AnySolved = false;

   /// Active binding (scope + choice) which is currently
   /// being attempted, helps to rewind state of the
   /// constraint system back to original before attempting
   /// next binding, if any.
   Optional<std::pair<std::unique_ptr<Scope>, typename P::Element>> ActiveChoice;

   BindingStep(ConstraintSystem &cs, P producer,
               SmallVectorImpl<Solution> &solutions)
      : SolverStep(cs, solutions), Producer(std::move(producer)) {}

public:
   StepResult take(bool prevFailed) override {
      while (auto choice = Producer()) {
         if (shouldSkip(*choice))
            continue;

         if (shouldStopAt(*choice))
            break;

         if (isDebugMode()) {
            auto &log = getDebugLogger();
            log << "(attempting ";
            choice->print(log, &CS.getAstContext().SourceMgr);
            log << '\n';
         }

         {
            auto scope = std::make_unique<Scope>(CS);
            if (attempt(*choice)) {
               ActiveChoice.emplace(std::move(scope), *choice);
               return suspend(std::make_unique<SplitterStep>(CS, Solutions));
            }
         }

         if (isDebugMode())
            getDebugLogger() << ")\n";

         // If this binding didn't match, let's check if we've attempted
         // enough bindings to stop, because some producers might need
         // to compute next step of bindings to try, which we'd want to avoid.
         if (shouldStopAfter(*choice))
            break;
      }

      return done(/*isSuccess=*/AnySolved);
   }

protected:
   /// Attempt to apply given binding choice to constraint system.
   /// This action is going to establish "active choice" of this step
   /// to point to a given choice.
   ///
   /// \param choice The choice to attempt.
   ///
   /// \return true if the choice has been accepted and system can be
   /// simplified further, false otherwise.
   virtual bool attempt(const typename P::Element &choice) = 0;

   /// Check whether attempting this choice could be avoided,
   /// which could speed-up solving.
   virtual bool shouldSkip(const typename P::Element &choice) const = 0;

   /// Check whether attempting binding choices should be stopped,
   /// because optimal solution has already been found.
   virtual bool shouldStopAt(const typename P::Element &choice) const = 0;

   /// Check whether attempting binding choices should be stopped,
   /// after current choice has been attempted, because optimal
   /// solution has already been found,
   virtual bool shouldStopAfter(const typename P::Element &choice) const {
      return false;
   }

   bool needsToComputeNext() const { return Producer.needsToComputeNext(); }

   ConstraintLocator *getLocator() const { return Producer.getLocator(); }
};

class TypeVariableStep final : public BindingStep<TypeVarBindingProducer> {
   using BindingContainer = ConstraintSystem::PotentialBindings;
   using Binding = ConstraintSystem::PotentialBinding;

   TypeVariableType *TypeVar;
   // A set of the initial bindings to consider, which is
   // also a source of follow-up "computed" bindings such
   // as supertypes, defaults etc.
   SmallVector<Binding, 4> InitialBindings;

   /// Indicates whether source of one of the previously
   /// attempted bindings was a literal constraint. This
   /// is useful for a performance optimization to stop
   /// attempting other bindings in certain conditions.
   bool SawFirstLiteralConstraint = false;

public:
   TypeVariableStep(ConstraintSystem &cs, BindingContainer &bindings,
                    SmallVectorImpl<Solution> &solutions)
      : BindingStep(cs, {cs, bindings}, solutions), TypeVar(bindings.TypeVar),
        InitialBindings(bindings.Bindings.begin(), bindings.Bindings.end()) {}

   void setup() override;

   StepResult resume(bool prevFailed) override;

   void print(llvm::raw_ostream &Out) override {
      PrintOptions PO;
      PO.PrintTypesForDebugging = true;
      Out << "TypeVariableStep for " << TypeVar->getString(PO) << " with #"
          << InitialBindings.size() << " initial bindings\n";
   }

protected:
   bool attempt(const TypeVariableBinding &choice) override;

   bool shouldSkip(const TypeVariableBinding &choice) const override {
      // If this is a defaultable binding and we have found solutions,
      // don't explore the default binding.
      return AnySolved && choice.isDefaultable();
   }

   /// Check whether attempting type variable binding choices should
   /// be stopped, because optimal solution has already been found.
   bool shouldStopAt(const TypeVariableBinding &choice) const override {
      // If we were able to solve this without considering
      // default literals, don't bother looking at default literals.
      return AnySolved && choice.hasDefaultedInterface() &&
             !SawFirstLiteralConstraint;
   }

   bool shouldStopAfter(const TypeVariableBinding &choice) const override {
      // If there has been at least one solution so far
      // at a current batch of bindings is done it's a
      // success because each new batch would be less
      // and less precise.
      return AnySolved && needsToComputeNext();
   }
};

class DisjunctionStep final : public BindingStep<DisjunctionChoiceProducer> {
   Constraint *Disjunction;
   SmallVector<Constraint *, 4> DisabledChoices;
   ConstraintList::iterator AfterDisjunction;

   Optional<Score> BestNonGenericScore;
   Optional<std::pair<Constraint *, Score>> LastSolvedChoice;

public:
   DisjunctionStep(ConstraintSystem &cs, Constraint *disjunction,
                   SmallVectorImpl<Solution> &solutions)
      : BindingStep(cs, {cs, disjunction}, solutions), Disjunction(disjunction),
        AfterDisjunction(erase(disjunction)) {
      assert(Disjunction->getKind() == ConstraintKind::Disjunction);
      pruneOverloadSet(Disjunction);
      ++cs.solverState->NumDisjunctions;
   }

   ~DisjunctionStep() override {
      // Rewind back any changes left after attempting last choice.
      ActiveChoice.reset();
      // Return disjunction constraint back to the system.
      restore(AfterDisjunction, Disjunction);
      // Re-enable previously disabled overload choices.
      for (auto *choice : DisabledChoices)
         choice->setEnabled();
   }

   StepResult resume(bool prevFailed) override;

   void print(llvm::raw_ostream &Out) override {
      Out << "DisjunctionStep for ";
      Disjunction->print(Out, &CS.getAstContext().SourceMgr);
      Out << '\n';
   }

private:
   bool shouldSkip(const DisjunctionChoice &choice) const override;

   /// Whether we should short-circuit a disjunction that already has a
   /// solution when we encounter the given choice.
   ///
   /// FIXME: This is performance hack, which should go away.
   ///
   /// \params choice The disjunction choice we are about to attempt.
   ///
   /// \returns true if disjunction step should be considered complete,
   ///          false otherwise.
   bool shouldStopAt(const DisjunctionChoice &choice) const override;
   bool shortCircuitDisjunctionAt(Constraint *currentChoice,
                                  Constraint *lastSuccessfulChoice) const;

   /// Attempt to apply given disjunction choice to constraint system.
   /// This action is going to establish "active choice" of this disjunction
   /// to point to a given choice.
   ///
   /// \param choice The choice to attempt.
   ///
   /// \return true if the choice has been accepted and system can be
   /// simplified further, false otherwise.
   bool attempt(const DisjunctionChoice &choice) override;

   // Check if selected disjunction has a representative
   // this might happen when there are multiple binary operators
   // chained together. If so, disable choices which differ
   // from currently selected representative.
   void pruneOverloadSet(Constraint *disjunction) {
      auto *choice = disjunction->getNestedConstraints().front();
      auto *typeVar = choice->getFirstType()->getAs<TypeVariableType>();
      if (!typeVar)
         return;

      auto *repr = typeVar->getImpl().getRepresentative(nullptr);
      if (!repr || repr == typeVar)
         return;

      for (auto *resolved = getResolvedOverloads(); resolved;
           resolved = resolved->Previous) {
         if (!resolved->BoundType->isEqual(repr))
            continue;

         auto &representative = resolved->Choice;
         if (!representative.isDecl())
            return;

         // Disable all of the overload choices which are different from
         // the one which is currently picked for representative.
         for (auto *constraint : disjunction->getNestedConstraints()) {
            auto choice = constraint->getOverloadChoice();
            if (!choice.isDecl() || choice.getDecl() == representative.getDecl())
               continue;

            constraint->setDisabled();
            DisabledChoices.push_back(constraint);
         }
         break;
      }
   };

   // Figure out which of the solutions has the smallest score.
   static Optional<Score> getBestScore(SmallVectorImpl<Solution> &solutions) {
      assert(!solutions.empty());
      Score bestScore = solutions.front().getFixedScore();
      if (solutions.size() == 1)
         return bestScore;

      for (unsigned i = 1, n = solutions.size(); i != n; ++i) {
         auto &score = solutions[i].getFixedScore();
         if (score < bestScore)
            bestScore = score;
      }
      return bestScore;
   }
};

} // end namespace constraints
} // end namespace polar

#endif // POLARPHP_SEMA_INTERNAL_CSSTEP_H
